(library (grid-gen GPU)
  (export GPU:make-code
          GPU:max-consts
  )
  (import (guile)
          (grid-gen main)
          (grid-gen vars)
          (grid-gen code)
          (grid-gen symath)
          (srfi srfi-1)
  )

(define GPU:max-comp-units 102400)

(define cur-locals (make-hash-table))
(define cur-reduce (make-hash-table))

(define rep-loops (lambda (cd)
  (if (and (list? cd) (not (null? cd)))
      (cond ((eq? (car cd) 'CODE_LOOP)
             (mk-loop cd)
            )
            ((eq? (car cd) 'CODE_BEGIN)
             (cons (car cd) (map rep-loops (cdr cd)))
            )
            ((eq? (car cd) 'OMP_LOCALS)
             (for-each
               (lambda (v) (hash-set! cur-locals (var-id v) v))
               (cdr cd)
             )
             cd
            )
            ((eq? (car cd) 'OMP_REDUCE)
             (let ((fn (obj-find-el cd 'OMP_FUNC))
                   (vr (car (reverse cd)))
                  )
               (hash-set! cur-reduce (var-id vr) (list fn vr))
               cd
             )
            )
            (#t cd)
      )
      cd
  )
))

(define cptr #f)

(define scan-vrs-ex (lambda (ex)
  (cond ((is-var? ex)
         (if (hash-ref cur-locals (var-id ex))
            '(()()()())
            `(()(,ex)()())
         )
        )
        ((or (is-op? ex) (is-fn? ex))
         (map (lambda (x) (apply append x))
              (apply zip (map scan-vrs-ex (cdr ex)))
         )
        )
        (#t '(()()()()))
  )  
))

(define scan-vrs (lambda (l)
  (if (and (list? l) (not (null? l)))
      (let ((op (car l)))
        (cond ((eq? op 'CODE_BEGIN)
               (map (lambda (x) (apply append x))
                    (apply zip (map scan-vrs (cdr l)))
               )
              )
              ((eq? op 'SET_PTR)
               (if (and cptr (not (equal? cptr l)))
                   (begin
                     (format #t "PTR: ~a\n" cptr)
                     (int-err "Multiple set-ptr is not supported!")
                   )
                   (set! cptr l)
               )
              '(()()()())
              )
              ((eq? op 'if)
               (map (lambda (x) (apply append x))
                    (zip (scan-vrs-ex (cadr l))
                         (scan-vrs (caddr l))
                         (if (not (null? (cadddr l)))
                             (scan-vrs (cadddr l))
                            '(()()()())
                         )
                    )
               )
              )
              ((eq? op 'CODE_WHILE)
               (int-err "While loop inside loop is not supported!")
              )
              ((eq? op 'CODE_LOOP)
               (int-err "Loop inside loop is not supported!")
              )
              ((eq? op 'CODE_NOP)
              '(()()()())
              )
              ((eq? op 'RAW)
              '(()()()())
              )
              ((eq? op 'OMP_LOCALS)
               (for-each
                 (lambda (v) (hash-set! cur-locals (var-id v) v))
                 (cdr l)
               )
              '(()()()())
              )
              ((eq? op 'OMP_REDUCE)
               (let ((fn (obj-find-el l 'OMP_FUNC))
                     (vr (car (reverse l)))
                    )
                 (hash-set! cur-reduce (var-id vr) (list fn vr))
                '(()()()())
               )
              )
              ((eq? op 'SET_VAL)
               (let ((vex (scan-vrs-ex (caddr l))))
                 (list (if (hash-ref cur-locals (var-id (cadr l)))
                          '()
                           (list (cadr l))
                       )
                       (cadr vex)
                      '()
                      '()
                 )
               )
              )
              (#t (begin
                    (format #t "OP: ~a\n" op)
                    (int-err "Unknown op!")
                  )
              )
        )
      )
      (begin
        (format #t "CD: ~a\n" l)
        (int-err "Invalid code")
      )
  )
))

(define cur-kern-id 0)

(define get-kern-id (lambda ()
  (let ((cid cur-kern-id))
    (set! cur-kern-id (+ cid 1))
    cid
  )
))

(define kernels (make-hash-table))

(define GPU:max-consts 0)

(define mk-loop (lambda (lp)
  (let* ((dms (cadr (obj-find-el lp 'LOOP_DIMS)))
         (seq (cadr (obj-find-el lp 'LOOP_SEQ)))
         (pre (obj-find-el lp 'LOOP_PRE_CODE))
         (cod (cadr (obj-find-el lp 'LOOP_CODE)))
        )
    (set! cptr #f)
    (letrec* ((vrs (map (lambda (x) (apply append x))
                        (zip (scan-vrs cod)
                             (if pre
                                (let ((vv (scan-vrs (cadr pre))))
                                 `(()(),(car vv),(cadr vv))
                                )
                               '(()()()())
                             )
                        )
                   )
              )
              (ovrs (let ((h (make-hash-table)))
                      (for-each 
                        (lambda (v) (hash-set! h (var-id v) v))
                        (append (car vrs) (caddr vrs))
                      )
                      h
                    )
              )
              (cnst (let ((h (make-hash-table)))
                      (for-each (lambda (v) (hash-set! h (var-id v) v))
                                (filter (lambda (v)
                                          (and (not (var-has-ind? v))
                                               (not (hash-ref ovrs 
                                                             (var-id v)
                                                             #f
                                                    )
                                               )
                                               (not (is-ptr? v))
                                          )
                                        )
                                        (apply append vrs)
                                )
                      )
                      h
                    )
              )
              (vrs2 (map (lambda (vl)
                           (filter (lambda (v)
                                     (and (not (hash-ref cnst (var-id v) #f))
                                          (not (hash-ref cur-reduce
                                                         (var-id v)
                                                         #f
                                               )
                                          )
                                     )
                                   )
                                   vl
                           )
                         )
                         vrs
                    )
              )
              (cnstl (hash-map->list (lambda (id vr) vr) cnst))
              (loc-vars (hash-map->list (lambda (i v) v) cur-locals))
              (reds (hash-map->list
                      (lambda (i rd) (cons 'OMP_REDUCE rd))
                      cur-reduce
                    )
              )
              (kern-id (get-kern-id))
              (lseq (car seq))
              (warp-size 
                (simp-ex
                  (cons '*
                        (map (lambda (x)
                               `(+ (- ,(cadr x) ,(car x)) 1)
                             )
                             (list-cut dms lseq)
                        )
                  )
                )
              )
              (warp-len 
                (let ((wl (list-ref dms lseq)))
                  (simp-ex `(+ (- ,(cadr wl) ,(car wl)) 1))
                )
              )
              (in-pre (cadddr vrs2))
              (in-loop (cadr vrs2))
              (out-pre 
                (filter
                  (lambda (v) (not (hash-ref cur-reduce (var-id v) #f)))
                  (caddr vrs2)
                )
              )
              (out-loop
                (filter
                  (lambda (v) (not (hash-ref cur-reduce (var-id v) #f)))
                  (car vrs2)
                )
              )
              (il-size (simp-ex `(+ ,(length in-pre)
                                    ,(- (length seq) 1)
                                     (* ,(length in-loop)
                                        ,warp-len
                                     )
                                 )
                       )
              )
              (ivr-size (simp-ex `(* ,warp-size ,il-size)))
              (ol-size (simp-ex `(+ ,(length out-pre)
                                     (* ,(length out-loop)
                                        ,warp-len
                                     )
                                 )
                       )
              )
              (ovr-size (simp-ex `(* ,warp-size ,ol-size)))
              (in-vrs (float-array (format #f "krn_in~a" kern-id)
                                   (list ivr-size)
                      )
              )
              (out-vrs (float-array (format #f "krn_out~a" kern-id)
                                    (list ovr-size)
                       )
              )

              (mk-lin (lambda (dims dms pt)
                (if (< (length dms) 1)
                    0
                    (let* ((d (list-ref dims (car dms)))
                           (dd `(+ (- ,(cadr d) ,(car d)) 1))
                          )
                         `(+ (- (CUR_PTR ,pt) ,(car d))
                             (* ,dd ,(mk-lin dims (cdr dms) (+ pt 1)))
                          )
                    )
                )
              ))
              (ivc (get-int-var))
              (cd0        (map (lambda (i)
                                 `(SET_VAL ,(set-ind
                                              in-vrs 
                                             `((+ (* ,i ,warp-size)
                                                 ,(mk-lin dms (cdr seq) 1)
                                                  1
                                               )
                                              )
                                            )
                                            (CUR_PTR ,(+ i 1))
                                  )
                                )
                                (iota1 (- (length seq) 1))
                           )
              )
             )
      (hash-set! kernels 
                 kern-id
                 (list in-pre
                       in-loop
                       out-pre
                       out-loop
                       cnstl
                       loc-vars
                       reds
                       lp
                       warp-size
                       ivr-size
                       ovr-size
                 )
      )
      (set! GPU:max-consts (max GPU:max-consts (length cnstl)))
     `(CODE_BEGIN
        (CODE_LOOP
          (LOOP_DIMS ,dms)
          (LOOP_SEQ  ,seq)
          (LOOP_PRE_CODE
           ,(cons 'CODE_BEGIN
              (let* ((cdl (map (lambda (x)
                                 `(SET_VAL
                                   ,(set-ind in-vrs
                                            `((+ (* ,(+ (car x) 
                                                        (- (length seq) 1)
                                                     )
                                                    ,warp-size
                                                 )
                                                ,ivc
                                                 1
                                              )  
                                             )
                                    )
                                   ,(cadr x)
                                  )
                               )
                               (zip (iota1 (length in-pre)) in-pre)
                          )
                     )
                    )
                (append
                  cd0
                 `((SET_VAL ,ivc ,(mk-lin dms (cdr seq) 1)))
                  (if cptr (list cptr) '())
                  cdl
                )
              )
            )
          )
          (LOOP_CODE
            (CODE_BEGIN
              (SET_VAL ,ivc
                       ,(simp-ex `(+ ,(mk-lin dms (cdr seq) 1)
                                      (+ (* (- (CUR_PTR 0)
                                              ,(car (list-ref dms (car seq)))
                                            )
                                            (* ,warp-size ,(length in-loop))
                                         )
                                         (* ,warp-size 
                                             (+ ,(length in-pre)
                                                ,(- (length seq) 1)
                                             )
                                         )
                                      )
                                     1
                                  )
                        )
              )
             ,(cons 'CODE_BEGIN
                    (append
                      (if cptr (list cptr) '())
                      (map (lambda (x)
                             `(SET_VAL
                                ,(set-ind in-vrs 
                                         `((+ ,ivc (* ,(car x) ,warp-size)))
                                 )
                                ,(cadr x)
                              )
                           )
                           (zip (iota1 (length in-loop)) in-loop)
                      )
                    )
              )
            )
          )
        )
        
        (CODE_CALL ,(format #f "gpu_run_kernel~a" kern-id)
                   ,(append (list in-vrs)
                            cnstl
                            (map car (map reverse reds))
                    )
        )

       ,(if (or (not (null? reds)) (> ovr-size 0))
           `(CODE_CALL ,(format #f "gpu_get_result~a" kern-id)
                       ,(append
                          (if (> ovr-size 0)
                            `(,out-vrs)
                            '()
                          )
                          (map (lambda (r) (car (reverse r))) reds)
                        )
            )
           '(CODE_NOP)
        )
        
       ,(if (or (not (null? out-pre))
                (not (null? out-loop))
            )
           `(CODE_LOOP
              (LOOP_DIMS ,dms)
              (LOOP_SEQ  ,seq)
              (LOOP_PRE_CODE
               ,(cons 'CODE_BEGIN
                  (let* ((cdl (map (lambda (x)
                                     `(SET_VAL
                                       ,(set-ind out-vrs
                                                `((+ (* ,(car x)
                                                        ,warp-size
                                                      )
                                                    ,(mk-lin dms (cdr seq) 1)
                                                     1
                                                  )  
                                                 )
                                        )
                                       ,(cadr x)
                                      )
                                   )
                                   (zip (iota1 (length out-pre)) out-pre)
                              )
                         )
                        )
                    (append
                      (if cptr (list cptr) '())
                      cdl
                    )
                  )
                )
              )
              (LOOP_CODE
                (CODE_BEGIN
                 ,(cons 'CODE_BEGIN
                      (append
                       `((SET_VAL ,ivc
                           ,(simp-ex `(+ ,(mk-lin dms (cdr seq) 1)
                                          (+ (* (- (CUR_PTR 0) 
                                                  ,(car (list-ref dms
                                                                  (car seq)
                                                        )
                                                   )
                                                )
                                                (* ,warp-size
                                                   ,(length out-loop)
                                                )
                                             )
                                             (* ,warp-size ,(length out-pre))
                                          )
                                          1
                                      )
                            )
                        ))
                        (if cptr (list cptr) '())
                        (map (lambda (x)
                               `(SET_VAL
                                  ,(cadr x)
                                  ,(set-ind out-vrs 
                                           `((+ ,ivc (* ,(car x) ,warp-size)))
                                   )
                                )
                             )
                             (zip (iota1 (length out-loop)) out-loop)
                        )
                      )
                  )
                )
              )
            )
           '(CODE_NOP)
        )
      ;;; ,lp
      )
    )
  )
))
  
(define do-wr (lambda (l)
  (set! cur-locals (make-hash-table))
  (set! cur-reduce (make-hash-table))
  (cons 'CODE_BEGIN (map rep-loops (force-code (cdr l))))
))

(define GPU:make-code (lambda (cdh)
  (int-err "GPU make-code function not found!")
))

(define init-wr (lambda (cod)
  (let ((cod1 (force-code (cdr cod))))
    (GPU:make-code kernels)
    (list (car cod)
          (cons 'CODE_BEGIN
                 (append '((CODE_CALL "gpu_init" ()))
                          (force-code cod1)
                         '((CODE_CALL "gpu_finish" ()))
                 )
          )
    )
  )
))
  
)

(import (grid-gen main))

(add-wrapper 12 'OMP_DO do-wr)
(add-wrapper 12 'RUN_CODE init-wr)
